# 本文件保存所有数据存储相关函数，可以保存至文件或者数据
# by yewentao
# v0.1
import pandas as pd
from os import path
import openpyxl
from openpyxl import Workbook
import _params

# 将股票列表数据保存至excel文件中，此方法为简单覆盖方式（如果已存在文件，将被覆盖）
def save_stock_list_to_file_simple(data, file_full_path = 'stock_list.xlsx'):
    if type(data) is pd.DataFrame :
        data.to_excel(file_full_path, index=False)
    else :
        raise TypeError('参数data数据类型必须是DataFrame!')

def save_stock_list_to_file_addition(data, field, file_full_path = 'stock_list.xlsx'):
    if type(data) is pd.DataFrame:
        # 如果文件已存在，则需要进行增量添加而不是覆盖
        if path.exists(file_full_path):
            # 注意，由于code是纯数字，默认读取到dataframe后，code变成int，例如：000001变成1，我们需要设置为str型
            stock_list = pd.read_excel(file_full_path, converters={field:str})
            # 对数据按照代码进行排序
            stock_list.sort_values(by = field)
            # 获取最大的代码值
            max_code = int(stock_list.iloc[len(stock_list)-1][field])
            # 获取最新数据中，代码比文件中已存在代码大的数据，用以增量保存
            add_data = data[data[field].apply(lambda x : int(x) > max_code)]
            #print("新增数据为：\n")
            #print(add_data)
            # 将新增数据增量写入excel
            workbook = openpyxl.load_workbook(file_full_path)
            table = workbook.get_sheet_by_name('Sheet1')
            for index, row in add_data.iterrows():
                table.append(list(row))
            workbook.save(file_full_path)
        else:
            data.to_excel(file_full_path, index=False)
    else :
        raise TypeError('参数data数据类型必须是DataFrame!')